"""
工作流节点模块 - 提供通用工作流节点
"""

from typing import Dict, Any, Callable, List
from langchain_openai import ChatOpenAI
import json
import re
import logging

from FDAS.config.settings import (
    OPENAI_API_KEY,
    OPENAI_BASE_URL,
    OPENAI_MODEL,
    OPENAI_TEMPERATURE
)

logger = logging.getLogger(__name__)

def parse_query_node(
    llm: ChatOpenAI,
    prompt_template: str,
    output_key: str = "parsed_query",
    error_handler: Callable = None
) -> Callable:
    """
    创建查询解析节点
    
    参数:
        llm: 大模型实例
        prompt_template: 提示模板
        output_key: 输出键名
        error_handler: 错误处理函数
        
    返回:
        节点函数
    """
    def node_func(state: Dict[str, Any]) -> Dict[str, Any]:
        try:
            # 获取用户查询
            user_query = state.get('user_query', '')
            
            # 替换提示模板中的变量
            prompt = prompt_template.format(user_query=user_query)
            
            # 调用大模型
            response = llm.invoke(prompt)
            
            # 解析响应
            result = _clean_json_response(response.content)
            
            # 返回结果
            return {
                **state,
                output_key: result
            }
        except Exception as e:
            logger.error(f"查询解析失败: {str(e)}")
            
            # 如果提供了错误处理函数，则调用它
            if error_handler:
                return error_handler(state, e)
            
            # 否则返回默认错误信息
            return {
                **state,
                output_key: {"error": str(e)}
            }
    
    return node_func

def retrieve_data_node(
    data_retriever: Callable,
    input_key: str,
    output_key: str = "data",
    error_handler: Callable = None
) -> Callable:
    """
    创建数据获取节点
    
    参数:
        data_retriever: 数据获取函数
        input_key: 输入键名
        output_key: 输出键名
        error_handler: 错误处理函数
        
    返回:
        节点函数
    """
    def node_func(state: Dict[str, Any]) -> Dict[str, Any]:
        try:
            # 获取输入数据
            input_data = state.get(input_key)
            
            # 调用数据获取函数
            result = data_retriever(input_data)
            
            # 返回结果
            return {
                **state,
                output_key: result
            }
        except Exception as e:
            logger.error(f"数据获取失败: {str(e)}")
            
            # 如果提供了错误处理函数，则调用它
            if error_handler:
                return error_handler(state, e)
            
            # 否则返回默认错误信息
            return {
                **state,
                output_key: {"error": str(e)}
            }
    
    return node_func

def analyze_data_node(
    llm: ChatOpenAI,
    prompt_template: str,
    data_key: str,
    output_key: str = "analysis_result",
    error_handler: Callable = None
) -> Callable:
    """
    创建数据分析节点 - 分析单一类型数据并生成总结
    
    参数:
        llm: 大模型实例
        prompt_template: 提示模板
        data_key: 数据键名
        output_key: 输出键名
        error_handler: 错误处理函数
        
    返回:
        节点函数
    """
    def node_func(state: Dict[str, Any]) -> Dict[str, Any]:
        try:
            # 获取数据
            data = state.get(data_key, {})
            
            # 检查数据是否有效
            if not data:
                raise ValueError("未获取到有效数据")
            
            if 'error' in data:
                raise ValueError(f"数据错误: {data['error']}")
            
            # 替换提示模板中的变量
            prompt = prompt_template.format(data=json.dumps(data, ensure_ascii=False), **state)
            
            # 调用大模型
            response = llm.invoke(prompt)
            
            # 返回结果 - 只保存分析总结，不保存原始数据
            return {
                **state,
                output_key: response.content
            }
        except Exception as e:
            logger.error(f"数据分析失败: {str(e)}")
            
            # 如果提供了错误处理函数，则调用它
            if error_handler:
                return error_handler(state, e)
            
            # 否则返回默认错误信息
            return {
                **state,
                output_key: _format_error_response(str(e))
            }
    
    return node_func

def analyze_specific_data_node(
    llm: ChatOpenAI,
    prompt_template: str,
    data_extractor: Callable,
    output_key: str,
    error_handler: Callable = None
) -> Callable:
    """
    创建特定数据分析节点 - 提取并分析特定类型的数据
    
    参数:
        llm: 大模型实例
        prompt_template: 提示模板
        data_extractor: 从状态中提取相关数据的函数
        output_key: 输出键名
        error_handler: 错误处理函数
        
    返回:
        节点函数
    """
    def node_func(state: Dict[str, Any]) -> Dict[str, Any]:
        try:
            # 使用提供的提取器获取特定数据
            data = data_extractor(state)
            
            if not data:
                raise ValueError(f"未获取到有效数据用于{output_key}分析")
            
            # 替换提示模板中的变量
            prompt = prompt_template.format(**data, **state)
            
            # 调用大模型
            response = llm.invoke(prompt)
            
            # 返回结果 - 添加到状态中，保留原有状态
            return {
                **state,
                output_key: response.content
            }
        except Exception as e:
            logger.error(f"{output_key}分析失败: {str(e)}")
            
            # 如果提供了错误处理函数，则调用它
            if error_handler:
                return error_handler(state, e)
            
            # 否则返回默认错误信息
            return {
                **state,
                output_key: _format_error_response(str(e))
            }
    
    return node_func

def synthesize_analysis_node(
    llm: ChatOpenAI,
    prompt_template: str,
    summary_keys: List[str],
    output_key: str = "final_analysis",
    error_handler: Callable = None
) -> Callable:
    """
    创建分析综合节点 - 整合多个分析总结生成最终报告
    
    参数:
        llm: 大模型实例
        prompt_template: 提示模板
        summary_keys: 要整合的分析总结键名列表
        output_key: 输出键名
        error_handler: 错误处理函数
        
    返回:
        节点函数
    """
    def node_func(state: Dict[str, Any]) -> Dict[str, Any]:
        try:
            # 收集所有分析总结
            summaries = {}
            for key in summary_keys:
                summary = state.get(key)
                if summary:
                    summaries[key] = summary
            
            if not summaries:
                raise ValueError("未找到有效的分析总结")
            
            # 构建提示
            prompt = prompt_template.format(
                summaries=json.dumps(summaries, ensure_ascii=False),
                **state
            )
            
            # 调用大模型生成最终分析
            response = llm.invoke(prompt)
            
            # 返回结果
            return {
                **state,
                output_key: response.content
            }
        except Exception as e:
            logger.error(f"综合分析失败: {str(e)}")
            
            # 如果提供了错误处理函数，则调用它
            if error_handler:
                return error_handler(state, e)
            
            # 否则返回默认错误信息
            return {
                **state,
                output_key: _format_error_response(str(e))
            }
    
    return node_func

def _clean_json_response(response: str) -> Dict:
    """
    清理并解析JSON响应
    
    参数:
        response: 大模型返回的JSON字符串
        
    返回:
        解析后的JSON字典
    """
    try:
        # 清理响应内容
        cleaned_content = re.sub(r'[\x00-\x1F\u200b-\u200d\ufe00-\ufe0f]', '', response)
        
        # 提取JSON - Python re模块不支持递归模式 (?R)，使用更简单的方法
        # 首先尝试直接解析整个内容
        try:
            return json.loads(cleaned_content)
        except json.JSONDecodeError:
            pass
        
        # 尝试提取第一个JSON对象
        json_start = cleaned_content.find('{')
        json_end = cleaned_content.rfind('}')
        
        if json_start != -1 and json_end != -1:
            json_str = cleaned_content[json_start:json_end+1]
            try:
                return json.loads(json_str)
            except json.JSONDecodeError:
                # 尝试修复常见格式问题
                fixed_json = re.sub(r',\s*}', '}', json_str)
                fixed_json = re.sub(r',\s*]', ']', fixed_json)
                return json.loads(fixed_json)
        
        # 如果上面的尝试都失败，使用简单的正则表达式
        json_match = re.search(r'\{[^{}]*\}', cleaned_content)
        if json_match:
            try:
                return json.loads(json_match.group())
            except json.JSONDecodeError:
                fixed_json = re.sub(r',\s*}', '}', json_match.group())
                fixed_json = re.sub(r',\s*]', ']', fixed_json)
                return json.loads(fixed_json)
        
        raise ValueError("未检测到有效JSON结构")
            
    except Exception as e:
        logger.error(f"JSON解析失败: {str(e)}")
        return {}

def _format_error_response(error_message: str) -> str:
    """
    格式化错误响应
    
    参数:
        error_message: 错误信息
        
    返回:
        格式化后的错误响应
    """
    return f"""
## 分析过程中出现错误

很抱歉，在分析过程中遇到了以下错误：

```
{error_message}
```

请尝试以下解决方案：

1. 检查输入的资产代码是否正确
2. 确保选择了正确的资产类型
3. 尝试选择不同的分析类型或时间范围
4. 如果问题持续存在，请联系系统管理员
    """